#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
微信聊天记录获取工具
使用 itchat 库获取微信群聊天的聊天记录
"""

import itchat
import os
import json
import time
from datetime import datetime
import pandas as pd
from itchat.content import *

class WeChatChatExporter:
    def __init__(self):
        self.chat_records = []
        self.export_dir = "exported_chats"
        
        # 创建导出目录
        if not os.path.exists(self.export_dir):
            os.makedirs(self.export_dir)
    
    def login(self):
        """登录微信"""
        print("正在启动微信登录...")
        print("请使用微信扫描二维码登录")
        
        # 启用热重载，避免重复登录
        itchat.auto_login(hotReload=True)
        print("登录成功！")
    
    def get_chatrooms(self):
        """获取所有群聊"""
        chatrooms = itchat.get_chatrooms()
        print(f"找到 {len(chatrooms)} 个群聊")
        
        # 显示群聊列表
        for i, room in enumerate(chatrooms):
            print(f"{i+1}. {room['NickName']} (成员数: {room['MemberCount']})")
        
        return chatrooms
    
    def select_chatroom(self, chatrooms):
        """选择要导出的群聊"""
        while True:
            try:
                choice = input("\n请输入要导出的群聊编号 (1-{}): ".format(len(chatrooms)))
                choice = int(choice) - 1
                if 0 <= choice < len(chatrooms):
                    selected_room = chatrooms[choice]
                    print(f"已选择: {selected_room['NickName']}")
                    return selected_room
                else:
                    print("无效的选择，请重新输入")
            except ValueError:
                print("请输入有效的数字")
    
    def get_chat_history(self, chatroom, limit=1000):
        """获取群聊聊天记录"""
        print(f"正在获取群聊 '{chatroom['NickName']}' 的聊天记录...")
        
        # 获取群聊的聊天记录
        messages = itchat.get_msg_list(toUserName=chatroom['UserName'], limit=limit)
        
        chat_records = []
        for msg in messages:
            record = {
                'timestamp': msg.get('CreateTime', 0),
                'datetime': datetime.fromtimestamp(msg.get('CreateTime', 0)).strftime('%Y-%m-%d %H:%M:%S'),
                'sender': msg.get('ActualNickName', '未知'),
                'content': self._extract_content(msg),
                'msg_type': msg.get('Type', '未知'),
                'msg_id': msg.get('MsgId', '')
            }
            chat_records.append(record)
        
        print(f"成功获取 {len(chat_records)} 条聊天记录")
        return chat_records
    
    def _extract_content(self, msg):
        """提取消息内容"""
        msg_type = msg.get('Type', '')
        
        if msg_type == 'Text':
            return msg.get('Text', '')
        elif msg_type == 'Picture':
            return '[图片]'
        elif msg_type == 'Voice':
            return '[语音]'
        elif msg_type == 'Video':
            return '[视频]'
        elif msg_type == 'File':
            return '[文件]'
        elif msg_type == 'Card':
            return '[名片]'
        elif msg_type == 'Location':
            return '[位置]'
        else:
            return '[其他类型消息]'
    
    def export_to_json(self, chat_records, chatroom_name):
        """导出聊天记录到JSON文件"""
        filename = f"{self.export_dir}/{chatroom_name}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
        
        with open(filename, 'w', encoding='utf-8') as f:
            json.dump(chat_records, f, ensure_ascii=False, indent=2)
        
        print(f"聊天记录已导出到: {filename}")
        return filename
    
    def export_to_excel(self, chat_records, chatroom_name):
        """导出聊天记录到Excel文件"""
        filename = f"{self.export_dir}/{chatroom_name}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.xlsx"
        
        df = pd.DataFrame(chat_records)
        df.to_excel(filename, index=False, engine='openpyxl')
        
        print(f"聊天记录已导出到: {filename}")
        return filename
    
    def export_to_txt(self, chat_records, chatroom_name):
        """导出聊天记录到TXT文件"""
        filename = f"{self.export_dir}/{chatroom_name}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.txt"
        
        with open(filename, 'w', encoding='utf-8') as f:
            f.write(f"群聊: {chatroom_name}\n")
            f.write(f"导出时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
            f.write("=" * 50 + "\n\n")
            
            for record in chat_records:
                f.write(f"[{record['datetime']}] {record['sender']}: {record['content']}\n")
        
        print(f"聊天记录已导出到: {filename}")
        return filename
    
    def run(self):
        """运行主程序"""
        try:
            # 登录微信
            self.login()
            
            # 获取群聊列表
            chatrooms = self.get_chatrooms()
            
            if not chatrooms:
                print("未找到任何群聊")
                return
            
            # 选择群聊
            selected_room = self.select_chatroom(chatrooms)
            
            # 获取聊天记录
            chat_records = self.get_chat_history(selected_room)
            
            if not chat_records:
                print("未获取到聊天记录")
                return
            
            # 选择导出格式
            print("\n请选择导出格式:")
            print("1. JSON格式")
            print("2. Excel格式")
            print("3. TXT格式")
            print("4. 全部格式")
            
            while True:
                try:
                    export_choice = input("请输入选择 (1-4): ")
                    export_choice = int(export_choice)
                    if 1 <= export_choice <= 4:
                        break
                    else:
                        print("无效的选择，请重新输入")
                except ValueError:
                    print("请输入有效的数字")
            
            # 导出聊天记录
            chatroom_name = selected_room['NickName']
            
            if export_choice == 1:
                self.export_to_json(chat_records, chatroom_name)
            elif export_choice == 2:
                self.export_to_excel(chat_records, chatroom_name)
            elif export_choice == 3:
                self.export_to_txt(chat_records, chatroom_name)
            elif export_choice == 4:
                self.export_to_json(chat_records, chatroom_name)
                self.export_to_excel(chat_records, chatroom_name)
                self.export_to_txt(chat_records, chatroom_name)
            
            print("\n导出完成！")
            
        except KeyboardInterrupt:
            print("\n程序被用户中断")
        except Exception as e:
            print(f"程序运行出错: {e}")
        finally:
            # 退出微信
            itchat.logout()
            print("已退出微信")

def main():
    """主函数"""
    print("=" * 50)
    print("微信聊天记录获取工具")
    print("=" * 50)
    print("注意: 使用此工具需要遵守相关法律法规")
    print("请确保您有权限获取相关聊天记录")
    print("=" * 50)
    
    # 确认用户同意
    confirm = input("您是否同意上述条款？(y/n): ").lower()
    if confirm != 'y':
        print("程序退出")
        return
    
    # 创建导出器并运行
    exporter = WeChatChatExporter()
    exporter.run()

if __name__ == "__main__":
    main()
